﻿using System;
using System.Collections.Generic;
using System.Text;
using System.Diagnostics;
using System.Xml.Serialization;
using System.Security;
using System.Collections.Specialized;
using System.IO;
using LucasCode.Win32;

namespace QuickTools.ProcessHelper
{
    [Flags]
    public enum ProcessType
    {
        AnyProcess = 0,
        UserProcess = 1,
        ForeignProcess = 2,
        GameProcess = 4,
    }
    public class ProcessInformation
    {
        public ProcessPriorityClass Priority = ProcessPriorityClass.Normal;
        public int CPUAffinity = int.MaxValue;
    }
    public class ProcessSetting
    {
        private string m_Name = "*";
        public string Name
        {
            get { return m_Name; }
            set { m_Name = value; }
        }

        private ProcessType m_ProcessType = ProcessType.UserProcess;
        public ProcessType ProcessType
        {
            get { return m_ProcessType; }
            set { m_ProcessType = value; }
        }
        /*[XmlIgnore]
        private Process m_Process = null;*/
        private bool m_Restore = true;
        public bool Restore
        {
            get { return m_Restore; }
            set { m_Restore = value; }
        }
        private ProcessInformation targetState = new ProcessInformation();
        //private ProcessInformation oldState = new ProcessInformation();
        public ProcessPriorityClass Priority
        {
            get { return targetState.Priority; }
            set { targetState.Priority = value; }
        }
        public int CPUAffinity
        {
            get { return targetState.CPUAffinity; }
            set
            {
                if (value > Math.Pow(Environment.ProcessorCount, 2))
                    return;
                targetState.CPUAffinity = value;
            }
        }
        private bool m_Kill = false;
        public bool Kill
        {
            get { return m_Kill; }
            set { m_Kill = value; }
        }
    }

    public class ProcessModule
    {
        private List<ProcessSetting> m_ProcessRules = new List<ProcessSetting>();
        public IEnumerable<ProcessSetting> ProcessRules
        {
            get
            {
                return m_ProcessRules;
            }
            set
            {
                m_ProcessRules.Clear();
                m_ProcessRules.AddRange(value);
            }
        }
        private Dictionary<int, ProcessInformation> m_OriginalSettings = new Dictionary<int, ProcessInformation>();
        private StringCollection m_KilledProcesses = new StringCollection();

        public ProcessModule()
        {
            ProcessSetting p = new ProcessSetting();
            p.CPUAffinity = 0;
            p.Priority = ProcessPriorityClass.Normal;
            p.ProcessType = ProcessType.GameProcess;
            p.Name = "*";
            m_ProcessRules.Add(p);

            ProcessSetting p2 = new ProcessSetting();
            p2.CPUAffinity = 1;
            p2.Priority = ProcessPriorityClass.BelowNormal;
            p2.ProcessType = ProcessType.UserProcess;
            p2.Name = "*";
            m_ProcessRules.Add(p2);
        }
        public ProcessSetting CreateNew()
        {
            ProcessSetting ps = new ProcessSetting();
            m_ProcessRules.Add(ps);
            return ps;
        }
        public bool Remove(ProcessSetting ps)
        {
            return m_ProcessRules.Remove(ps);
        }
        private Process m_GameProcess;

        public void GameEnded()
        {
            foreach (KeyValuePair<int, ProcessInformation> kvp in m_OriginalSettings)
            {
                try
                {
                    Process p = Process.GetProcessById(kvp.Key);
                    if (p == null)
                        continue;
                    p.ProcessorAffinity = new IntPtr(kvp.Value.CPUAffinity);
                    p.PriorityClass = kvp.Value.Priority;
                    p.Dispose();
                }
                catch (Exception ex)
                {
                    LogError("Could not restore process: " + ex.Message);
                }
            }
            m_OriginalSettings.Clear();
            foreach (string s in m_KilledProcesses)
            {
                try
                {
                    using (Process p = Process.Start(s))
                    {
                        LogInfo(String.Format("Restarting process {0}", p.ProcessName));
                    }
                }
                catch (System.Exception ex)
                {
                    LogError(String.Format("Error restarting process {0}: {1}", Path.GetFileName(s), ex.Message));
                }
            }
            m_KilledProcesses.Clear();
        }

        private unsafe static bool GlobCompare(string glob, string path)
        {
            fixed (char* pGlob = glob, pPath = path)
            {
                int pGlobInc = 0;
                int pPathInc = 0;

                int mp = 0;
                int cp = 0;

                while ((*(pPath + pPathInc) != 0) && (*(pGlob + pGlobInc) != '*'))
                {
                    if ((*(pGlob + pGlobInc) != *(pPath + pPathInc)) && (*(pGlob + pGlobInc) != '?'))
                    {
                        return false;
                    }
                    pGlobInc++;
                    pPathInc++;
                }

                while (*(pPath + pPathInc) != 0)
                {
                    if (*(pGlob + pGlobInc) == '*')
                    {
                        if (0 == *(pGlob + ++pGlobInc))
                        {
                            return true;
                        }
                        mp = pGlobInc;
                        cp = pPathInc + 1;
                    }
                    else if ((*(pGlob + pGlobInc) == *(pPath + pPathInc)) || (*(pGlob + pGlobInc) == '?'))
                    {
                        pGlobInc++;
                        pPathInc++;
                    }
                    else
                    {
                        pGlobInc = mp;
                        pPathInc = cp++;
                    }
                }

                while (*(pGlob + pGlobInc) == '*')
                {
                    pGlobInc++;
                }
                return (0 == *(pGlob + pGlobInc));
            }
        }

        private ProcessType GetTypeOfProcess(Process p)
        {
            if (m_GameProcess != null && p.Id == m_GameProcess.Id)
                return ProcessType.GameProcess;

            string sid = Win32Process.GetProcessOwnerSID(p);
            if (System.Security.Principal.WindowsIdentity.GetCurrent().User.Value == sid)
                return ProcessType.UserProcess;
            return ProcessType.ForeignProcess;
        }
        private void Log(string format, params object[] args)
        {
            Trace.WriteLine(String.Format(format, args));
        }
        protected void LogInfo(string text, params object[] args)
        {
            Log(text, args);
        }
        protected void LogError(string text, params object[] args)
        {
            Log(text, args);
        }
        public void GameStarted(/*IReminderInterface sender, GameStartedEventArgs args*/)
        {
            //m_GameProcess = args.Game.Process;
            m_OriginalSettings.Clear();
            Process[] procs = Process.GetProcesses();
            foreach (Process p in procs)
            {
                // skip Idle and System
                if (p.Id == 0 || p.Id == 4)
                    continue;
                foreach (ProcessSetting ps in m_ProcessRules)
                {
                    if (GlobCompare(ps.Name.ToLower(), p.ProcessName.ToLower()) && ((ps.ProcessType & GetTypeOfProcess(p)) == ps.ProcessType))
                    {
                        try
                        {
                            if (ps.Kill)
                            {
                                LogInfo(String.Format("Killing process {0}", p.ProcessName));
                                p.Kill();
                                if (ps.Restore)
                                    m_KilledProcesses.Add(p.MainModule.FileName);
                            }
                            else
                            {
                                ProcessInformation pi = new ProcessInformation();
                                pi.CPUAffinity = p.ProcessorAffinity.ToInt32();
                                pi.Priority = p.PriorityClass;

                                p.ProcessorAffinity = new IntPtr(ps.CPUAffinity);
                                p.PriorityClass = ps.Priority;
                                if (ps.Restore)
                                    m_OriginalSettings.Add(p.Id, pi);
                            }
                        }
                        catch (Exception ex)
                        {
                            LogError(String.Format("Could not apply process settings for {0} ({1}) to {2}: {3}", ps.Name, Enum.GetName(typeof(ProcessType), ps.ProcessType), p.ProcessName, ex.Message));
                        }
                    }
                }
                p.Dispose();
            }
        }
    }
}
